//
// Created by vivi on 11/01/2018.
//

#include "HairCore.h"
#include <numeric>

namespace XRwy {
	using namespace std;
	using namespace Eigen;

	HairStictionCageSolver::HairStictionCageSolver(const HairBaseSimulatorConfiguration & conf):
		HairStictionSolver(),
		stiction_rate(conf.cage_stiction_rate), _1_minus_stiction_rate(1 - conf.cage_stiction_rate) {}

	void HairStictionCageSolver::solve(HairBaseSimulator &simulator, const HairSimulatorData &data,
	                               const Eigen::Affine3f &transformation, const float t) {
		for (auto & particle_view : simulator.particles)
			if (particle_view.local_id != 0)
				particle_view.pos() = _1_minus_stiction_rate * particle_view.pos()
				                      + stiction_rate * (transformation * particle_view.rest_pos);
	}

	HairStictionDensityFieldSolver::HairStictionDensityFieldSolver(const HairBaseSimulatorConfiguration & conf):
		HairStictionSolver(),
		stiction_rate(conf.density_stiction_rate),
		_1_minus_stiction_rate(1 - conf.density_stiction_rate),
		cell_count(conf.density_cell_count),
		vertex_count(cell_count + 1),
		vertex_count_2(vertex_count * vertex_count),
		vertex_count_3(vertex_count * vertex_count_2),
		vertexes(new HairStictionDensityFieldVertex[vertex_count_3]),
		position_infos(nullptr)
	{}

	HairStictionDensityFieldSolver::~HairStictionDensityFieldSolver() {
		delete [] vertexes;
		if (position_infos != nullptr)
			delete [] position_infos;
	}

	void HairStictionDensityFieldSolver::solve(HairBaseSimulator &simulator, const HairSimulatorData &data,
	                                           const Eigen::Affine3f &transformation, const float t) {
		//allocate the space if needed
		if (position_infos == nullptr)
			position_infos = new HairStictionDensityFieldParticlePositionInfo[simulator.particles.size()];

		//bounding box should needs to be computed
		min_x = min_y = min_z = std::numeric_limits<float>::max();
		max_x = max_y = max_z = std::numeric_limits<float>::min();

		for (const auto & particle_view : simulator.particles) {
			const auto & pos = particle_view.pos();
			min_x = std::min(min_x, pos(0));
			min_y = std::min(min_y, pos(1));
			min_z = std::min(min_z, pos(2));
			max_x = std::max(max_x, pos(0));
			max_y = std::max(max_y, pos(1));
			max_z = std::max(max_z, pos(2));
		}

		//extend a little bit, so the boundary error won't occur
		dx = (max_x - min_x) / 256.0f;
		dy = (max_y - min_y) / 256.0f;
		dz = (max_z - min_z) / 256.0f;
		max_x += dx;
		min_x -= dx;
		max_y += dy;
		min_y -= dy;
		max_z += dz;
		min_z -= dz;

		dx = (max_x - min_x) / cell_count;
		dy = (max_y - min_y) / cell_count;
		dz = (max_z - min_z) / cell_count;

		//initialize all the vertexes
		for (size_t i = 0; i < vertex_count_3; ++i) {
			vertexes[i].vel = Eigen::Vector3f::Zero();
			vertexes[i].density = 0;
		}

		//update the position infos
		for (size_t k = 0; k < simulator.particles.size(); ++k) {
			const auto & particle_view = simulator.particles[k];
			const auto & pos = particle_view.pos();
			auto & _ = position_infos[k];

			_.x = (pos(0) - min_x) / dx;
			_.y = (pos(1) - min_y) / dy;
			_.z = (pos(2) - min_z) / dz;
			//since the largest x will yield a value of vertex_count, we will limit it to the vertex_count - 1
			_.cx = static_cast<size_t>(_.x);
			_.cy = static_cast<size_t>(_.y);
			_.cz = static_cast<size_t>(_.z);

			HairCoreMath::trilinear_interpolation_values(_.x - _.cx, _.y - _.cy, _.z - _.cz, _.interpolation_weights);

			HairStictionDensityFieldVertexPointer s = vertexes + _.cx * vertex_count_2 + _.cy * vertex_count + _.cz;
			for (size_t i = 0; i < 8; ++i) {
				_.surrounding_vertexes[i] = s
				                            + (i & 4 ? vertex_count_2 : 0)
				                            + (i & 2 ? vertex_count : 0)
				                            + (i & 1 ? 1 : 0);

				_.surrounding_vertexes[i]->density += _.interpolation_weights[i];
				_.surrounding_vertexes[i]->vel += _.interpolation_weights[i] * particle_view.vel();
			}
		}

		//reupdate the particles
		for (size_t k = 0; k < simulator.particles.size(); ++k) {
			auto & particle_view = simulator.particles[k];
			const auto & _ = position_infos[k];

			Vector3f v_grid = Vector3f::Zero();
			for (size_t i = 0; i < 8; ++i)
				v_grid += _.interpolation_weights[i] * _.surrounding_vertexes[i]->vel / _.surrounding_vertexes[i]->density;

			particle_view.vel() = _1_minus_stiction_rate * particle_view.vel()
			                      + stiction_rate * v_grid;
		}
 	}

	HairStictionSpringSolver::HairStictionSpringSolver(const HairBaseSimulatorConfiguration & conf):
		HairStictionSolver(),
		distance(conf.spring_distance),
		tolerance(conf.spring_tolerance),
		acc(conf.spring_acc)
	{}

	HairStictionSpringSolver::~HairStictionSpringSolver() {
		for (size_t i = 0; i < 4; ++i)
			if (orders[i] != nullptr)
				delete [] (orders[i]);
		if (links != nullptr)
			delete [] links;
	}

	void HairStictionSpringSolver::solve(HairBaseSimulator &simulator, const HairSimulatorData &data,
	                                     const Eigen::Affine3f &transformation, const float t) {
		//if it has not been initialized
		if (links == nullptr) {
			links = new list<size_t>[simulator.particles.size()] ();
			for (size_t k = 0; k < 4; ++k) {
				orders[k] = new ParticleView* [simulator.particles.size()];
				for (size_t i = 0; i < simulator.particles.size(); ++i)
					orders[k][i] = &(simulator.particles[i]);
			}
		}
	}
}
